from random import randint

from utils.raw_obj import RawObject as Obj
from algo.sort.heap import Heap
from utils.generator import Generator
from utils.log import Log
from utils.profile import Profile


class MedianNumber(Log):
    # ref: https://www.bilibili.com/video/BV1nF411y7rD/?spm_id_from=333.788.recommend_more_video.-1&vd_source=ad559e41b561d48861f70d0187e51e02
    # 00:50:59
    def __init__(self):
        super(MedianNumber, self).__init__()
        self.max_heap = Heap(max_heap=True)
        self.min_heap = Heap(max_heap=False)

    @Profile(log=True)
    def add(self, val):
        v = val
        if not isinstance(val, Obj):
            # 必须转换成Obj对象,因为小根堆只支持Obj对象
            v = Obj(val)
        if (len(self.max_heap) == 0) or (self.max_heap.top() >= val):
            self.max_heap.insert(v)
        else:
            self.min_heap.insert(v)
        self._balance()

    def _balance(self):
        max_size = len(self.max_heap)
        min_size = len(self.min_heap)
        if abs(max_size - min_size) == 2:
            if max_size > min_size:
                self.min_heap.insert(self.max_heap.remove(0))
            else:
                self.max_heap.insert(self.min_heap.remove(0))

    def get(self):
        max_size = len(self.max_heap)
        min_size = len(self.min_heap)
        if max_size == min_size:
            if max_size > 0:
                return (int(self.max_heap.top()) + int(self.min_heap.top())) / 2
        else:
            return int(self.max_heap.top()) if max_size > min_size else int(self.min_heap.top())


if __name__ == '__main__':
    g = Generator()
    mn = MedianNumber()
    for _ in range(5):
        v = g.single_val()
        mn.add(v)
    print(mn.max_heap.get_arr())
    print(mn.min_heap.get_arr())
    print(mn.get())
